package main

import (
	"context"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/sync/semaphore"
)

type PortScanner struct {
	ip   string
	lock *semaphore.Weighted
}

func (ps PortScanner) Scan(port int, timeout time.Duration) error {
	address := fmt.Sprintf("%s:%d", ps.ip, port)
	conn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		return err
	}
	conn.Close()
	return nil
}

func (ps PortScanner) Start(from, to int, timeout time.Duration) {
	wg := sync.WaitGroup{}
	defer wg.Wait()
	fmt.Println("Scanning ip", ps.ip)
	if to == 0 {
		to = from + 1
	}
	for port := from; port < to; port++ {
		ps.lock.Acquire(context.TODO(), 1)
		wg.Add(1)
		go func(port int) {
			defer ps.lock.Release(1)
			defer wg.Done()
			start := time.Now()
			err := ps.Scan(port, timeout)
			if err == nil {
				fmt.Println("Port", port, "open, time", time.Since(start).Milliseconds(), "ms")
			} else {
				fmt.Println("Port", port, "closed")
			}
		}(port)
	}
}

func main() {
	helps := []string{}
	helps = append(helps, "\r\ncommand error, eg:")
	helps = append(helps, "1. portscanner 127.0.0.1:3000")
	helps = append(helps, "2. portscanner 127.0.0.1:1/65536")
	tips := strings.Join(helps, "\r\n")
	if len(os.Args) != 2 {
		panic(tips)
	}

	address := os.Args[1]
	s := strings.Split(address, ":")
	if len(s) != 2 {
		panic(tips)
	}

	ip := s[0]
	s = strings.Split(s[1], "/")
	var from, to int
	if len(s) > 0 {
		from, _ = strconv.Atoi(s[0])
	}
	if len(s) > 1 {
		to, _ = strconv.Atoi(s[1])
	}

	ps := PortScanner{
		ip:   ip,
		lock: semaphore.NewWeighted(100),
	}
	ps.Start(from, to, 1*time.Second)
}
